// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#pragma once

#include <pollux/common/memory/byte_stream.h>

namespace kumo::pollux::exec {
    /// Corresponds to Presto SerializedPage, i.e. a container for serialize vectors
    /// in Presto wire format.
    class SerializedPage {
    public:
        /// Construct from IOBuf chain.
        explicit SerializedPage(
            std::unique_ptr<melon::IOBuf> iobuf,
            std::function<void(melon::IOBuf &)> onDestructionCb = nullptr,
            std::optional<int64_t> numRows = std::nullopt);

        ~SerializedPage();

        /// Returns the size of the serialized data in bytes.
        uint64_t size() const {
            return iobufBytes_;
        }

        std::optional<int64_t> numRows() const {
            return numRows_;
        }

        /// Makes 'input' ready for deserializing 'this' with
        /// VectorStreamGroup::read().
        std::unique_ptr<ByteInputStream> prepareStreamForDeserialize();

        std::unique_ptr<melon::IOBuf> getIOBuf() const {
            return iobuf_->clone();
        }

    private:
        static int64_t chainBytes(melon::IOBuf &iobuf) {
            int64_t size = 0;
            for (auto &range: iobuf) {
                size += range.size();
            }
            return size;
        }

        // Buffers containing the serialized data. The memory is owned by 'iobuf_'.
        std::vector<ByteRange> ranges_;

        // IOBuf holding the data in 'ranges_.
        std::unique_ptr<melon::IOBuf> iobuf_;

        // Number of payload bytes in 'iobuf_'.
        const int64_t iobufBytes_;

        // Number of payload rows, if provided.
        const std::optional<int64_t> numRows_;

        // Callback that will be called on destruction of the SerializedPage,
        // primarily used to free externally allocated memory backing melon::IOBuf
        // from caller. Caller is responsible to pass in proper cleanup logic to
        // prevent any memory leak.
        std::function<void(melon::IOBuf &)> onDestructionCb_;
    };

    /// Queue of results retrieved from source. Owned by shared_ptr by
    /// Exchange and client threads and registered callbacks waiting
    /// for input.
    class ExchangeQueue {
    public:
        explicit ExchangeQueue(
            int32_t numberOfConsumers,
            uint64_t minOutputBatchBytes)
            : numberOfConsumers_{numberOfConsumers},
              minOutputBatchBytes_{minOutputBatchBytes} {
            POLLUX_CHECK_GE(numberOfConsumers, 1);
        }

        ~ExchangeQueue() {
            clearAllPromises();
        }

        std::mutex &mutex() {
            return mutex_;
        }

        bool empty() const {
            return queue_.empty();
        }

        /// Enqueues 'page' to the queue. One random promise(top of promise queue)
        /// associated with the future that is waiting for the data from the queue is
        /// returned in 'promises' if 'page' is not nullptr. When 'page' is nullptr
        /// and the queue is completed serving data, all left over promises will be
        /// returned in 'promises'. When 'page' is nullptr and the queue is not
        /// completed serving data, no 'promises' will be added and returned.
        void enqueueLocked(
            std::unique_ptr<SerializedPage> &&page,
            std::vector<ContinuePromise> &promises);

        /// If data is permanently not available, e.g. the source cannot be
        /// contacted, this registers an error message and causes the reading
        /// Exchanges to throw with the message.
        void setError(const std::string &error);

        /// Returns pages of data.
        ///
        /// Returns empty list if no data is available. If data is still expected,
        /// sets 'atEnd' to false and 'future' to a Future that will complete when
        /// data arrives. If no more data is expected, sets 'atEnd' to true. Returns
        /// at least one page if data is available. If multiple pages are available,
        /// returns as many pages as fit within 'maxBytes', but no fewer than one.
        /// Calling this method with 'maxBytes' of 1 returns at most one page.
        ///
        /// The data may be compressed, in which case 'maxBytes' applies to compressed
        /// size.
        std::vector<std::unique_ptr<SerializedPage> > dequeueLocked(
            int consumerId,
            uint32_t maxBytes,
            bool *atEnd,
            ContinueFuture *future,
            ContinuePromise *stalePromise);

        /// Returns the total bytes held by SerializedPages in 'this'.
        int64_t totalBytes() const {
            return totalBytes_;
        }

        /// Returns the maximum value of total bytes.
        uint64_t peakBytes() const {
            return peakBytes_;
        }

        /// Returns total number of pages received from all sources.
        uint64_t receivedPages() const {
            return receivedPages_;
        }

        /// Returns an average size of received pages. Returns 0 if hasn't received
        /// any pages yet.
        uint64_t averageReceivedPageBytes() const {
            return receivedPages_ > 0 ? receivedBytes_ / receivedPages_ : 0;
        }

        void addSourceLocked() {
            POLLUX_CHECK(!noMoreSources_, "addSource called after noMoreSources");
            numSources_++;
        }

        void noMoreSources();

        void close();

    private:
        std::vector<ContinuePromise> closeLocked() {
            queue_.clear();
            return clearAllPromisesLocked();
        }

        std::vector<ContinuePromise> checkCompleteLocked() {
            if (noMoreSources_ && numCompleted_ == numSources_) {
                atEnd_ = true;
                return clearAllPromisesLocked();
            }
            return {};
        }

        void addPromiseLocked(
            int consumerId,
            ContinueFuture *future,
            ContinuePromise *stalePromise);

        void clearAllPromises() {
            std::vector<ContinuePromise> promises; {
                std::lock_guard<std::mutex> l(mutex_);
                promises = clearAllPromisesLocked();
            }
            clearPromises(promises);
        }

        std::vector<ContinuePromise> clearAllPromisesLocked() {
            std::vector<ContinuePromise> promises(promises_.size());
            auto it = promises_.begin();
            while (it != promises_.end()) {
                promises.push_back(std::move(it->second));
                it = promises_.erase(it);
            }
            POLLUX_CHECK(promises_.empty());
            return promises;
        }

        static void clearPromises(std::vector<ContinuePromise> &promises) {
            for (auto &promise: promises) {
                promise.setValue();
            }
        }

        int64_t minOutputBatchBytesLocked() const;

        const int32_t numberOfConsumers_;
        const uint64_t minOutputBatchBytes_;

        int numCompleted_{0};
        int numSources_{0};
        bool noMoreSources_{false};
        bool atEnd_{false};

        std::mutex mutex_;
        std::deque<std::unique_ptr<SerializedPage> > queue_;
        // The map from consumer id to the waiting promise
        std::unordered_map<int, ContinuePromise> promises_;

        // When set, all promises will be realized and the next dequeue will
        // throw an exception with this message.
        std::string error_;
        // Total size of SerializedPages in queue.
        int64_t totalBytes_{0};
        // Number of SerializedPages received.
        int64_t receivedPages_{0};
        // Total size of SerializedPages received. Used to calculate an average
        // expected size.
        int64_t receivedBytes_{0};
        // Maximum value of totalBytes_.
        int64_t peakBytes_{0};
    };
} // namespace kumo::pollux::exec
